use std::hash::Hash;
use std::fmt::Debug;
use indexmap::IndexMap;
use serde::Serialize;
use serde_json::Value;

/// 树节点特征，任何想要构建树的结构都需要实现这个特征
pub trait TreeNode {
    type Id: Eq + Hash + Clone + Debug;
    
    /// 获取节点ID
    fn id(&self) -> Self::Id;
    
    /// 获取父节点ID，返回None表示根节点
    fn parent_id(&self) -> Option<Self::Id>;

    fn get_parent_ids(&self) -> Vec<Self::Id>;
}

/// 可序列化的树节点包装器
#[derive(Debug, Serialize)]
#[serde(bound = "T: Serialize, T::Id: Serialize")]
pub struct TreeNodeWrapper<'a, T: TreeNode> {
    /// 原始数据引用（会被 flatten，要求 T: Serialize）
    #[serde(flatten)]
    pub data: &'a T,
    /// 子节点
    pub children: Vec<TreeNodeWrapper<'a, T>>,
    /// 父级 ID 列表（从最近父级到最顶层）
    pub parent_ids: Vec<T::Id>,
}

impl<'a, T> TreeNodeWrapper<'a, T>
where
    T: TreeNode + Serialize,
    T::Id: Serialize,
{
    pub fn new(data: &'a T) -> Self {
        Self {
            data,
            children: Vec::new(),
            parent_ids: Vec::new(),
        }
    }
    
    /// 递归打印树结构
    pub fn print_tree(&self, depth: usize) {
        let indent = "  ".repeat(depth);
        println!("{}{:?} (parents: {:?})", indent, self.data.id(), self.parent_ids);
        
        for child in &self.children {
            child.print_tree(depth + 1);
        }
    }
    
    /// 查找特定ID的节点
    pub fn find_node(&self, id: &T::Id) -> Option<&TreeNodeWrapper<'a, T>> {
        if &self.data.id() == id {
            return Some(self);
        }
        
        for child in &self.children {
            if let Some(found) = child.find_node(id) {
                return Some(found);
            }
        }
        None
    }
    
    /// 获取所有后代节点（包括自身）
    pub fn descendants(&self) -> Vec<&TreeNodeWrapper<'a, T>> {
        let mut result = vec![self];
        for child in &self.children {
            result.extend(child.descendants());
        }
        result
    }

    /// 前序遍历
    pub fn pre_order_traversal(&self) -> Vec<&TreeNodeWrapper<'a, T>> {
        let mut result = vec![self];
        for child in &self.children {
            result.extend(child.pre_order_traversal());
        }
        result
    }
    
    /// 后序遍历
    pub fn post_order_traversal(&self) -> Vec<&TreeNodeWrapper<'a, T>> {
        let mut result = Vec::new();
        for child in &self.children {
            result.extend(child.post_order_traversal());
        }
        result.push(self);
        result
    }
    
    /// 检查是否有子节点
    pub fn has_children(&self) -> bool {
        !self.children.is_empty()
    }
    
    /// 获取子节点数量
    pub fn child_count(&self) -> usize {
        self.children.len()
    }
}

/// 包装器树结构
#[derive(Debug, Serialize)]
#[serde(bound = "T: Serialize, T::Id: Serialize")]
pub struct WrappedTree<'a, T: TreeNode> {
    pub roots: Vec<TreeNodeWrapper<'a, T>>,
    #[serde(skip)]
    pub total_nodes: usize,
}

impl<'a, T> WrappedTree<'a, T>
where
    T: TreeNode + Serialize,
    T::Id: Serialize,
{
    pub fn new(roots: Vec<TreeNodeWrapper<'a, T>>) -> Self {
        let total_nodes = roots.iter().map(|root| root.descendants().len()).sum();
        Self { roots, total_nodes }
    }
    
    /// 打印整棵树
    pub fn print(&self) {
        for root in &self.roots {
            root.print_tree(0);
        }
    }
    
    /// 查找特定ID的节点
    pub fn find_node(&self, id: &T::Id) -> Option<&TreeNodeWrapper<'a, T>> {
        for root in &self.roots {
            if let Some(found) = root.find_node(id) {
                return Some(found);
            }
        }
        None
    }
    
    /// 获取所有节点
    pub fn all_nodes(&self) -> Vec<&TreeNodeWrapper<'a, T>> {
        let mut result = Vec::new();
        for root in &self.roots {
            result.extend(root.descendants());
        }
        result
    }
    
    pub fn is_empty(&self) -> bool {
        self.roots.is_empty()
    }
    
    pub fn root_count(&self) -> usize {
        self.roots.len()
    }

    pub fn filter<F>(&self, predicate: F) -> Vec<&TreeNodeWrapper<'a, T>>
    where
        F: Fn(&T) -> bool,
    {
        self.all_nodes()
            .into_iter()
            .filter(|node| predicate(node.data))
            .collect()
    }
    
    pub fn map<U, F>(&self, mapper: F) -> Vec<U>
    where
        F: Fn(&TreeNodeWrapper<'a, T>) -> U,
    {
        self.all_nodes()
            .into_iter()
            .map(mapper)
            .collect()
    }
    
    pub fn depth(&self) -> usize {
        self.roots
            .iter()
            .map(|root| self.node_depth(root))
            .max()
            .unwrap_or(0)
    }
    
    fn node_depth(&self, node: &TreeNodeWrapper<'a, T>) -> usize {
        if node.children.is_empty() {
            1
        } else {
            1 + node.children
                .iter()
                .map(|child| self.node_depth(child))
                .max()
                .unwrap_or(0)
        }
    }
    
    pub fn to_json(&self) -> Result<String, serde_json::Error> {
        serde_json::to_string_pretty(self)
    }

    pub fn to_value(&self) -> Result<Value, serde_json::Error> {
        serde_json::to_value(&self.roots)
    }
    
    pub fn total_nodes(&self) -> usize {
        self.total_nodes
    }
}

/// 构建包装器树的通用函数
pub fn build_wrapped_tree<'a, T>(nodes: &'a [T]) -> WrappedTree<'a, T>
where
    T: TreeNode + Serialize,
    T::Id: Serialize,
{
    let mut node_refs: IndexMap<T::Id, TreeNodeWrapper<'a, T>> = IndexMap::new();
    let mut children_map: IndexMap<Option<T::Id>, Vec<T::Id>> = IndexMap::new();
    
    for node in nodes {
        let id = node.id();
        let parent_id = node.parent_id();
        node_refs.insert(id.clone(), TreeNodeWrapper::new(node));
        children_map.entry(parent_id.clone())
            .or_insert_with(Vec::new)
            .push(id);
    }
    
    let mut roots = Vec::new();

    // 根节点
    if let Some(root_ids) = children_map.get(&None) {
        for root_id in root_ids {
            if let Some(root_node) = build_wrapped_subtree(root_id, &children_map, &mut node_refs, Vec::new()) {
                roots.push(root_node);
            }
        }
    }

    // 处理孤立节点
    let mut orphan_root_ids = Vec::new();
    for (id, node_wrapper) in node_refs.iter() {
        if let Some(pid) = node_wrapper.data.parent_id() {
            if !node_refs.contains_key(&pid) {
                orphan_root_ids.push(id.clone());
            }
        }
    }

    for id in orphan_root_ids {
        let node = node_refs.get(&id);
        let parent_chain;
        if node.is_none() {
            parent_chain = Vec::new();
        } else {
            parent_chain = node.unwrap().data.get_parent_ids();
        }
        println!("parent_chain {:?}", parent_chain);
        if let Some(root_node) = build_wrapped_subtree(&id, &children_map, &mut node_refs, parent_chain) {
            roots.push(root_node);
        }
    }

    WrappedTree::new(roots)
}

/// 递归构建包装器子树（带 parent_ids）
fn build_wrapped_subtree<'a, T>(
    node_id: &T::Id,
    children_map: &IndexMap<Option<T::Id>, Vec<T::Id>>,
    node_refs: &mut IndexMap<T::Id, TreeNodeWrapper<'a, T>>,
    parent_chain: Vec<T::Id>,
) -> Option<TreeNodeWrapper<'a, T>>
where
    T: TreeNode + Serialize,
    T::Id: Serialize,
{
    let mut node = node_refs.swap_remove(node_id)?;
    node.parent_ids = parent_chain.clone();

    if let Some(child_ids) = children_map.get(&Some(node_id.clone())) {
        for child_id in child_ids {
            let mut new_chain = parent_chain.clone();
            new_chain.push(node_id.clone());
            if let Some(child_node) = build_wrapped_subtree(child_id, children_map, node_refs, new_chain) {
                node.children.push(child_node);
            }
        }
    }

    Some(node)
}
